#!/bin/bash
# =================================================================
# 使用 LWG + Buffer Gate 配置训练 MOTIP（SMD 数据集，8-GPU）
# 建议在 tmux 或类似工具中运行此脚本
# =================================================================

set -euo pipefail

# --- 基础配置 ---
PROJECT_DIR="/home/liuyonghui/code/MA-MOTIP"
cd "${PROJECT_DIR}" || exit

OUTPUTS_DIR="${OUTPUTS_DIR:-${PROJECT_DIR}/outputs/motip_smd_wandb}"
EXP_NAME="${EXP_NAME:-motip_smd_lwg}"

export WANDB_ENTITY="yves-det"
export WANDB_PROJECT="MOTIP"
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export WANDB_BASE_URL=https://api.bandw.top

DATA_ROOT="${DATA_ROOT:-/home/liuyonghui/datasets}"
PRETRAIN_PATH="${PRETRAIN_PATH:-${PROJECT_DIR}/pretrains/r50_deformable_detr_coco_dancetrack.pth}"
CONFIG_PATH="${CONFIG_PATH:-${PROJECT_DIR}/configs/r50_deformable_detr_motip_smd_lwg.yaml}"

# LWG / Buffer Gate 设置
LWG_MODEL_PATH="${LWG_MODEL_PATH:-${PROJECT_DIR}/pretrains/lwg_smd.pt}"
BUFFER_GATE_MODEL_PATH="${BUFFER_GATE_MODEL_PATH:-${PROJECT_DIR}/pretrains/buffer_gate_smd.pt}"
LWG_MAIN_TH="${LWG_MAIN_TH:-0.60}"
LWG_BUF_TH="${LWG_BUF_TH:-0.40}"
BUFFER_CAPACITY="${BUFFER_CAPACITY:-10}"
BUFFER_PROMOTE_STEPS="${BUFFER_PROMOTE_STEPS:-3}"

NUM_GPUS="${NUM_GPUS:-8}"

# --- 环境变量 ---
export NVI_NOTIFY_IGNORE_TASK=1
export LD_PRELOAD=/opt/lib/fakeintel/libfakeintel.so

echo "Activating conda environment..."
source /opt/miniconda3/etc/profile.d/conda.sh
conda activate MOTIP

# 日志
LOG_DIR="${PROJECT_DIR}/logs"
mkdir -p "${LOG_DIR}"
LOG_FILE="${LOG_DIR}/${EXP_NAME}_$(date +%Y%m%d_%H%M%S).log"

echo "项目目录: ${PROJECT_DIR}"
echo "数据集目录: ${DATA_ROOT}"
echo "输出目录: ${OUTPUTS_DIR}"
echo "日志文件: ${LOG_FILE}"
echo "GPU数量: ${NUM_GPUS}"

# --- 检查 ---
if [ ! -f "${PRETRAIN_PATH}" ]; then
    echo "错误: 未找到预训练模型 -> ${PRETRAIN_PATH}"
    exit 1
fi
if [ ! -d "${DATA_ROOT}" ]; then
    echo "错误: 数据集目录不存在 -> ${DATA_ROOT}"
    exit 1
fi

# --- 启动训练 ---
echo "开始训练..."
accelerate launch --num_processes="${NUM_GPUS}" train.py \
    --config-path "${CONFIG_PATH}" \
    --data-root "${DATA_ROOT}" \
    --detr-pretrain "${PRETRAIN_PATH}" \
    --outputs-dir "${OUTPUTS_DIR}" \
    --exp-name "${EXP_NAME}" \
    --use-wandb True \
    --use-lwg True \
    --lwg-model-path "${LWG_MODEL_PATH}" \
    --lwg-main-th "${LWG_MAIN_TH}" \
    --lwg-buf-th "${LWG_BUF_TH}" \
    --buffer-capacity "${BUFFER_CAPACITY}" \
    --buffer-promote-steps "${BUFFER_PROMOTE_STEPS}" \
    2>&1 | tee "${LOG_FILE}"

echo "训练完成。日志已保存到 ${LOG_FILE}"
